package nl.ViewerServer.ServerManager.JobManager;

import java.util.HashMap;
import java.util.Iterator;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

import nl.ViewerServer.Application;
import nl.ViewerServer.ServerManager.JobManager.Strategy.AssignmentStrategy;
import nl.ViewerServer.ServerManager.JobManager.Strategy.DefaultStrategy;
import nl.ViewerServer.ServerManager.JobManager.Strategy.LocalityStrategy;
import nl.ViewerServer.ServerManager.Paraview.Worker;
import nl.ViewerServer.SessionManager.Objects.Session;

import org.apache.log4j.Logger;

public class JobManager extends ThreadPoolExecutor {
   private static Logger         log     = Logger.getLogger( JobManager.class );
   // hashmap of sessionId's mapped to Worker threads
   private HashMap<Long, Worker> workers = new HashMap<Long, Worker>();

   private AssignmentStrategy    strategy;

   public JobManager() {
      super( Application.config.poolSize, Application.config.poolSize, 1, TimeUnit.MINUTES,
            new LinkedBlockingQueue<Runnable>() );

      // if the config defines a list of paraview nodes use locality strategy
      if( Application.config.paraViewMachines != null && Application.config.paraViewMachines.size() > 0 ) {
         strategy = new LocalityStrategy();
         log.info( "Using LocalityStrategy" );
      } else {
         strategy = new DefaultStrategy();
         log.info( "Using DefaultStrategy" );
      }
      log.info( "JobManager ThreadPool created" );
   }

   public void shutdown() {
      log.info( "JobManager shutting down" );
      // shutdown all workers
      Iterator<Worker> i = workers.values().iterator();
      while( i.hasNext() ) {
         Worker w = i.next();
         if( w != null ) {
            w.shutdown();
         }
      }
   }

   public synchronized ParaViewThreadAssignment assignParaviewThreads( int requestedThreads ) {
      return strategy.assignParaViewThreads( requestedThreads );
   }

   public synchronized void reclaimParaviewThreadAssignment( ParaViewThreadAssignment assignment ) {
      strategy.reclaimParaviewThreadAssignment( assignment );
   }

   public synchronized void submitSession( Session s ) {
      Worker w = workers.get( s.sessionId );
      if( w != null ) { // worker was already assigned, wake it if necessary
         w.interrupt();
         return;
      }
      s.setStatus( Session.Status.QUEUED );
      w = new Worker( s );
      workers.put( s.sessionId, w );
      /* Future<Session> future = */
      submit( Executors.callable( w, s ) );
   }

   /**
    * Called from the sessionManager to request a worker associated with a given session to be closed
    * 
    * @param s
    */

   public synchronized void closeWorker( Session s ) {
      Worker w = workers.get( s.sessionId );
      if( w == null )
         return;

      remove( w );
      purge();

      // start shutdown of the worker in the background (without waiting at all)
      w.asynchronousShutdown( 0 );
   }

   /**
    * Called from the sessionManager to request a worker to stop applying a timeout for new jobs
    * 
    * @param s
    */

   public synchronized void disableWorkerTimeout( Session s ) {
      Worker w = workers.get( s.sessionId );
      if( w == null )
         return;

      w.setTimeoutDisabled( true );
   }

   protected void beforeExecute( Thread t, Runnable r ) {
      super.beforeExecute( t, r );
   }

   protected void afterExecute( Runnable r, Throwable t ) {
      try {
         Session s = ( (Future<Session>)r ).get();

         // save the worker object to check for errors
         Worker w = workers.get( s.sessionId );

         // do some administration: remove the worker from the jobManager
         if( workers.remove( s.sessionId ) == null ) {
            log.warn( "Worker could not be removed from queue (session id " + s.sessionId + ")" );
         }

         // if the worker stopped because of an error, the associated session should be re-submitted
         if( w != null && w.getLastError() != null ) {
            if( s.getStatus() == Session.Status.FAILED ) {
               log.error( "Session failed: " + s.getFailureReason() );
            } else if( s.getRetryCount() < Application.config.sessionRetryCount ) {
               // the session should be restarted with a new worker!
               this.submitSession( s );
            } else {
               s.setStatusFailed( "This visualization failed too often, session was closed" );
               log.warn( "Worker for session " + s.sessionId + " failed too often, session is not re-submitted" );
            }
         }

      }
      catch( Exception e ) {
         if( log != null )
            log.warn( "Unexpected error while running worker", e );

      }
      super.afterExecute( r, t );
   }

}
